package com.feidee.fd.sml.algorithm.feature

import com.feidee.fd.sml.algorithm.component.feature.{InteractEncoder, InteractEncoderParam}
import com.feidee.fd.sml.algorithm.util.{TestingDataGenerator, ToolClass}
import org.apache.spark.ml.feature.VectorAssembler
import org.apache.spark.ml.linalg.Vectors
import org.scalatest.FunSuite

/**
  * @Author tangjinyuan
  * @Date 2019/3/28 14:32
  * @Description
  * @Reviewer
  */
class InteractEncoderSuite extends FunSuite {
  val paramStr: String =
    """
      |{
      |	"inputCol": "id1,vec1,vec2",
      |	"outputCol": "interactedCol"
      |}
    """.stripMargin

  val ie = new InteractEncoder()
  val param: InteractEncoderParam = ie.parseParam(new ToolClass().encrypt(paramStr))

  test("model parameter") {
    assert("interactedCol".equals(param.outputCol) && param.inputCol.split(",").length == 3 )
  }


  test("model transformation") {
    val compareSeq = Seq(
      (Vectors.dense(8.0,4.0,5.0,16.0,8.0,10.0,24.0,12.0,15.0),""),
      (Vectors.dense(56.0,72.0,64.0,42.0,54.0,48.0,112.0,144.0,128.0),""),
      (Vectors.dense(36.0,54.0,108.0,6.0,9.0,18.0,54.0,81.0,162.0),""),
      (Vectors.dense(360.0,160.0,200.0,288.0,128.0,160.0,216.0,96.0,120.0),""),
      (Vectors.dense(450.0,315.0,135.0,100.0,70.0,30.0,350.0,245.0,105.0),""),
      (Vectors.dense(12.0,48.0,24.0,12.0,48.0,24.0,48.0,192.0,96.0),"")
    )
    val comparaData = TestingDataGenerator.spark.createDataFrame(compareSeq).toDF("interactedCol","b")

    comparaData.show()
    val df = TestingDataGenerator.spark.createDataFrame(Seq(
      (1, 1, 2, 3, 8, 4, 5),
      (2, 4, 3, 8, 7, 9, 8),
      (3, 6, 1, 9, 2, 3, 6),
      (4, 10, 8, 6, 9, 4, 5),
      (5, 9, 2, 7, 10, 7, 3),
      (6, 1, 1, 4, 2, 8, 4)
    )).toDF("id1", "id2", "id3", "id4", "id5", "id6", "id7")

    val assembler1 = new VectorAssembler().
      setInputCols(Array("id2", "id3", "id4")).
      setOutputCol("vec1")

    val assembled1 = assembler1.transform(df)

    val assembler2 = new VectorAssembler().
      setInputCols(Array("id5", "id6", "id7")).
      setOutputCol("vec2")

    val assembled2 = assembler2.transform(assembled1).select("id1", "vec1", "vec2")

    val res = ie.train(param, assembled2).transform(assembled2)
    println("res:")
    res.show()
    comparaData.select("interactedCol").show()

    comparaData.select("interactedCol").printSchema()
    res.select("interactedCol").printSchema()
    assert(comparaData.select("interactedCol").except(res.select("interactedCol")).count() == 0)
  }

}

